{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    div#notebook-container    { width: 95%; }\n",
       "    div#menubar-container     { width: 65%; }\n",
       "    div#maintoolbar-container { width: 99%; }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 19, 05:40:15] Starting\n",
      "[Feb 19, 05:40:15] Learning rate: 0.001\n",
      "[Feb 19, 05:40:17] Number of learnable parameters: 33299701\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Anaconda3\\lib\\site-packages\\torch\\serialization.py:250: UserWarning: Couldn't retrieve source code for container of type Duet. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "C:\\Anaconda3\\lib\\site-packages\\torch\\serialization.py:250: UserWarning: Couldn't retrieve source code for container of type Flatten. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 19, 06:17:23] model:1, epoch:1, loss:0.23297182732494548\n",
      "[Feb 19, 07:45:08] model:1, mrr:0.24274774037871585\n",
      "[Feb 19, 07:45:09] Number of learnable parameters: 33299701\n",
      "[Feb 19, 08:22:15] model:2, epoch:1, loss:0.23497310453967657\n",
      "[Feb 19, 09:48:21] model:2, mrr:0.2463984090630236\n",
      "[Feb 19, 09:48:21] Number of learnable parameters: 33299701\n",
      "[Feb 19, 10:25:29] model:3, epoch:1, loss:0.2385119370155735\n",
      "[Feb 19, 11:51:33] model:3, mrr:0.24835987167077864\n",
      "[Feb 19, 11:51:33] Number of learnable parameters: 33299701\n",
      "[Feb 19, 12:28:43] model:4, epoch:1, loss:0.23345718482596567\n",
      "[Feb 19, 13:54:39] model:4, mrr:0.2500117582102976\n",
      "[Feb 19, 13:54:39] Number of learnable parameters: 33299701\n",
      "[Feb 19, 14:31:52] model:5, epoch:1, loss:0.23169705012696795\n",
      "[Feb 19, 15:58:05] model:5, mrr:0.25276113451392535\n",
      "[Feb 19, 15:58:05] Number of learnable parameters: 33299701\n",
      "[Feb 19, 16:34:58] model:6, epoch:1, loss:0.3315703275220585\n",
      "[Feb 19, 18:01:14] model:6, mrr:0.252444401274203\n",
      "[Feb 19, 18:01:14] Number of learnable parameters: 33299701\n",
      "[Feb 19, 18:37:30] model:7, epoch:1, loss:0.693147613608744\n",
      "[Feb 19, 20:03:36] model:7, mrr:0.252444401274203\n",
      "[Feb 19, 20:03:37] Number of learnable parameters: 33299701\n",
      "[Feb 19, 20:40:46] model:8, epoch:1, loss:0.2314062070217915\n",
      "[Feb 19, 22:06:52] model:8, mrr:0.25146449361307593\n",
      "[Feb 19, 22:13:14] Finished\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import sys\n",
    "import os\n",
    "import os.path\n",
    "import csv\n",
    "import re\n",
    "import random\n",
    "import datetime\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import torch.optim as optim\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "display(HTML(data=\"\"\"\n",
    "<style>\n",
    "    div#notebook-container    { width: 95%; }\n",
    "    div#menubar-container     { width: 65%; }\n",
    "    div#maintoolbar-container { width: 99%; }\n",
    "</style>\n",
    "\"\"\"))\n",
    "\n",
    "def print_message(s):\n",
    "    print(\"[{}] {}\".format(datetime.datetime.utcnow().strftime(\"%b %d, %H:%M:%S\"), s), flush=True)\n",
    "\n",
    "\n",
    "class DataReader:\n",
    "\n",
    "    def __init__(self, data_file, num_meta_cols, multi_pass):\n",
    "        self.num_meta_cols                  = num_meta_cols\n",
    "        self.multi_pass                     = multi_pass\n",
    "        self.regex_drop_char                = re.compile('[^a-z0-9\\s]+')\n",
    "        self.regex_multi_space              = re.compile('\\s+')\n",
    "        self.__load_vocab()\n",
    "        self.__load_idfs()\n",
    "        self.__init_data(data_file)\n",
    "        self.__allocate_minibatch()\n",
    "\n",
    "    def __tokenize(self, s, max_terms):\n",
    "        return self.regex_multi_space.sub(' ', self.regex_drop_char.sub(' ', s.lower())).strip().split()[:max_terms]\n",
    "\n",
    "    def __load_vocab(self):\n",
    "        global VOCAB_SIZE\n",
    "        self.vocab                          = {}\n",
    "        with open(DATA_FILE_VOCAB, mode='r', encoding=\"utf-8\") as f:\n",
    "            reader                          = csv.reader(f, delimiter='\\t')\n",
    "            for row in reader:\n",
    "                self.vocab[row[0]]          = int(row[1])\n",
    "        VOCAB_SIZE                          = len(self.vocab) + 1\n",
    "        embeddings                          = np.zeros((VOCAB_SIZE, NUM_HIDDEN_NODES), dtype=np.float32)\n",
    "        with open(DATA_EMBEDDINGS, mode='r', encoding=\"utf-8\") as f:\n",
    "            for line in f:\n",
    "                cols                                                    = line.split()\n",
    "                idx                                                     = self.vocab.get(cols[0], 0)\n",
    "                if idx > 0:\n",
    "                    for i in range(NUM_HIDDEN_NODES):\n",
    "                        embeddings[idx, i]                              = float(cols[i + 1])\n",
    "        self.pre_trained_embeddings                                     = torch.tensor(embeddings)\n",
    "        \n",
    "    def __load_idfs(self):\n",
    "        self.idfs                           = {}\n",
    "        with open(DATA_FILE_IDFS, mode='r', encoding=\"utf-8\") as f:\n",
    "            reader                          = csv.reader(f, delimiter='\\t')\n",
    "            for row in reader:\n",
    "                self.idfs[row[0]]           = float(row[1])\n",
    "\n",
    "    def __init_data(self, file_name):\n",
    "        self.reader                         = open(file_name, mode='r', encoding=\"utf-8\")\n",
    "        self.num_docs                       = len(self.reader.readline().split('\\t')) - self.num_meta_cols - 1\n",
    "        self.reader.seek(0)\n",
    "    \n",
    "    def __allocate_minibatch(self):\n",
    "        self.features                       = {}\n",
    "        if ARCH_TYPE != 1:\n",
    "            self.features['local']          = [] \n",
    "        if ARCH_TYPE > 0:\n",
    "            self.features['dist_q']         = np.zeros((MB_SIZE, MAX_QUERY_TERMS), dtype=np.int64)\n",
    "            self.features['mask_q']         = np.zeros((MB_SIZE, MAX_QUERY_TERMS, 1), dtype=np.float32)\n",
    "            self.features['dist_d']         = []\n",
    "            self.features['mask_d']         = []\n",
    "        for i in range(self.num_docs):\n",
    "            if ARCH_TYPE != 1:\n",
    "                self.features['local'].append(np.zeros((MB_SIZE, MAX_DOC_TERMS, MAX_QUERY_TERMS), dtype=np.float32))\n",
    "            if ARCH_TYPE > 0:\n",
    "                self.features['dist_d'].append(np.zeros((MB_SIZE, MAX_DOC_TERMS), dtype=np.int64))\n",
    "                self.features['mask_d'].append(np.zeros((MB_SIZE, MAX_DOC_TERMS, 1), dtype=np.float32))\n",
    "        self.features['labels']             = np.zeros((MB_SIZE), dtype=np.int64)\n",
    "        self.features['meta']               = []\n",
    "        \n",
    "    def __clear_minibatch(self):\n",
    "        if ARCH_TYPE > 0:\n",
    "            self.features['dist_q'].fill(np.int64(0))\n",
    "            self.features['mask_q'].fill(np.float32(0))\n",
    "        for i in range(self.num_docs):\n",
    "            if ARCH_TYPE != 1:\n",
    "                self.features['local'][i].fill(np.float32(0))\n",
    "            if ARCH_TYPE > 0:\n",
    "                self.features['dist_d'][i].fill(np.int64(0))\n",
    "                self.features['mask_d'][i].fill(np.float32(0))\n",
    "        self.features['meta'].clear()\n",
    "\n",
    "    def get_minibatch(self):\n",
    "        self.__clear_minibatch()\n",
    "        for i in range(MB_SIZE):\n",
    "            row                             = self.reader.readline()\n",
    "            if row == '':\n",
    "                if self.multi_pass:\n",
    "                    self.reader.seek(0)\n",
    "                    row                     = self.reader.readline()\n",
    "                else:\n",
    "                    break\n",
    "            cols                            = row.split('\\t')\n",
    "            q                               = self.__tokenize(cols[self.num_meta_cols], MAX_QUERY_TERMS)\n",
    "            ds                              = [self.__tokenize(cols[self.num_meta_cols + i + 1], MAX_DOC_TERMS) for i in range(self.num_docs)]\n",
    "            if ARCH_TYPE != 1:\n",
    "                for d in range(self.num_docs):\n",
    "                    for j in range(len(ds[d])):\n",
    "                        for k in range(len(q)):\n",
    "                            if ds[d][j] == q[k]:\n",
    "                                self.features['local'][d][i, j, k] = self.idfs[q[k]]\n",
    "            if ARCH_TYPE > 0:\n",
    "                for j in range(self.num_docs + 1):\n",
    "                    terms = q if j == 0 else ds[j - 1]\n",
    "                    for t in range(len(terms)):\n",
    "                        term = self.vocab.get(terms[t], 0)\n",
    "                        if j == 0:\n",
    "                            self.features['dist_q'][i, t] = term\n",
    "                            self.features['mask_q'][i, t, 0] = 1\n",
    "                        else:\n",
    "                            self.features['dist_d'][j - 1][i, t] = term\n",
    "                            self.features['mask_d'][j - 1][i, t, 0] = 1\n",
    "            self.features['meta'].append(tuple(cols[:self.num_meta_cols]))\n",
    "        return self.features\n",
    "    \n",
    "    def reset(self):\n",
    "        self.reader.seek(0)\n",
    "\n",
    "class Flatten(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Flatten, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x.contiguous().view(x.size(0), -1)\n",
    "\n",
    "class Duet(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(Duet, self).__init__()\n",
    "        self.embed                  = nn.Embedding(VOCAB_SIZE, NUM_HIDDEN_NODES)\n",
    "        self.embed.weight           = nn.Parameter(READER_TRAIN.pre_trained_embeddings, requires_grad=True)\n",
    "        self.duet_local             = nn.Sequential(nn.Conv1d(MAX_DOC_TERMS, NUM_HIDDEN_NODES, kernel_size=1),\n",
    "                                        nn.ReLU(),\n",
    "                                        Flatten(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES * MAX_QUERY_TERMS, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE))\n",
    "        self.duet_dist_q            = nn.Sequential(nn.Conv1d(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES, kernel_size=3),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.MaxPool1d(POOLING_KERNEL_WIDTH_QUERY),\n",
    "                                        Flatten(),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU()\n",
    "                                        )\n",
    "        self.duet_dist_d            = nn.Sequential(nn.Conv1d(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES, kernel_size=3),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.MaxPool1d(POOLING_KERNEL_WIDTH_DOC, stride=1),\n",
    "                                        nn.Conv1d(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES, kernel_size=1),\n",
    "                                        nn.ReLU()\n",
    "                                        )\n",
    "        self.duet_dist              = nn.Sequential(Flatten(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES * NUM_POOLING_WINDOWS_DOC, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE))\n",
    "        self.duet_comb              = nn.Sequential(nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES, NUM_HIDDEN_NODES),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Dropout(p=DROPOUT_RATE),\n",
    "                                        nn.Linear(NUM_HIDDEN_NODES, 1),\n",
    "                                        nn.ReLU())\n",
    "        self.scale                  = torch.tensor([0.1], requires_grad=False).to(DEVICE)\n",
    "\n",
    "    def forward(self, x_local, x_dist_q, x_dist_d, x_mask_q, x_mask_d):\n",
    "        if ARCH_TYPE != 1:\n",
    "            h_local                 = self.duet_local(x_local)\n",
    "        if ARCH_TYPE > 0:\n",
    "            h_dist_q                = self.duet_dist_q((self.embed(x_dist_q) * x_mask_q).permute(0, 2, 1))\n",
    "            h_dist_d                = self.duet_dist_d((self.embed(x_dist_d) * x_mask_d).permute(0, 2, 1))\n",
    "            h_dist                  = self.duet_dist(h_dist_q.unsqueeze(-1)*h_dist_d)\n",
    "        y_score                     = self.duet_comb((h_local + h_dist) if ARCH_TYPE == 2 else (h_dist if ARCH_TYPE == 1 else h_local)) * self.scale\n",
    "        return y_score\n",
    "    \n",
    "    def parameter_count(self):\n",
    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
    "\n",
    "DEVICE                              = torch.device(\"cuda:0\")    # torch.device(\"cpu\"), if you want to run on CPU instead\n",
    "ARCH_TYPE                           = 2\n",
    "MAX_QUERY_TERMS                     = 20\n",
    "MAX_DOC_TERMS                       = 200\n",
    "NUM_HIDDEN_NODES                    = 300\n",
    "TERM_WINDOW_SIZE                    = 3\n",
    "POOLING_KERNEL_WIDTH_QUERY          = MAX_QUERY_TERMS - TERM_WINDOW_SIZE + 1 # 20 - 3 + 1 = 18\n",
    "POOLING_KERNEL_WIDTH_DOC            = 100\n",
    "NUM_POOLING_WINDOWS_DOC             = (MAX_DOC_TERMS - TERM_WINDOW_SIZE + 1) - POOLING_KERNEL_WIDTH_DOC + 1 # (200 - 3 + 1) - 100 + 1 = 99\n",
    "VOCAB_SIZE                          = 0\n",
    "DROPOUT_RATE                        = 0.5\n",
    "MB_SIZE                             = 1024\n",
    "EPOCH_SIZE                          = 1024\n",
    "NUM_EPOCHS                          = 1\n",
    "NUM_ENSEMBLES                       = 8\n",
    "LEARNING_RATE                       = 0.001\n",
    "DATA_DIR                            = 'data\\\\'\n",
    "DATA_FILE_VOCAB                     = os.path.join(DATA_DIR, \"word-vocab-small.tsv\")\n",
    "DATA_EMBEDDINGS                     = os.path.join(DATA_DIR, \"glove.6B\\\\glove.6B.{}d.txt\".format(NUM_HIDDEN_NODES))\n",
    "DATA_FILE_IDFS                      = os.path.join(DATA_DIR, \"idfnew.norm.tsv\")\n",
    "DATA_FILE_TRAIN                     = os.path.join(DATA_DIR, \"triples.train.full.tsv\")\n",
    "DATA_FILE_DEV                       = os.path.join(DATA_DIR, \"top1000.dev.tsv\")\n",
    "DATA_FILE_EVAL                      = os.path.join(DATA_DIR, \"top1000.eval.tsv\")\n",
    "QRELS_DEV                           = os.path.join(DATA_DIR, \"qrels.dev.tsv\")\n",
    "DATA_FILE_OUT_DEV                   = os.path.join(DATA_DIR, \"output.dev.tsv\")\n",
    "DATA_FILE_OUT_EVAL                  = os.path.join(DATA_DIR, \"output.eval.tsv\")\n",
    "MODEL_FILE                          = os.path.join(DATA_DIR, \"duet.ens{}.ep{}.dnn\")\n",
    "READER_TRAIN                        = DataReader(DATA_FILE_TRAIN, 0, True)\n",
    "READER_DEV                          = DataReader(DATA_FILE_DEV, 2, False)\n",
    "READER_EVAL                         = DataReader(DATA_FILE_EVAL, 2, False)\n",
    "\n",
    "qrels                               = {}\n",
    "with open(QRELS_DEV, mode='r', encoding=\"utf-8\") as f:\n",
    "    reader                          = csv.reader(f, delimiter='\\t')\n",
    "    for row in reader:\n",
    "        qid                         = int(row[0])\n",
    "        did                         = int(row[2])\n",
    "        if qid not in qrels:\n",
    "            qrels[qid]              = []\n",
    "        qrels[qid].append(did)\n",
    "\n",
    "res_dev                             = {}\n",
    "res_eval                            = {}\n",
    "\n",
    "print_message('Starting')\n",
    "print_message('Learning rate: {}'.format(LEARNING_RATE))\n",
    "for ens_idx in range(NUM_ENSEMBLES):\n",
    "    torch.manual_seed(ens_idx + 1)\n",
    "    net                             = Duet()\n",
    "    net                             = net.to(DEVICE)\n",
    "    criterion                       = nn.CrossEntropyLoss()\n",
    "    optimizer                       = optim.Adam(net.parameters(), lr=LEARNING_RATE)\n",
    "    print_message('Number of learnable parameters: {}'.format(net.parameter_count()))\n",
    "    for ep_idx in range(NUM_EPOCHS):\n",
    "        train_loss                  = 0.0\n",
    "        net.train()\n",
    "        for mb_idx in range(EPOCH_SIZE):\n",
    "            features                = READER_TRAIN.get_minibatch()\n",
    "            if ARCH_TYPE == 0:\n",
    "                out                 = torch.cat(tuple([net(torch.from_numpy(features['local'][i]).to(DEVICE), None, None) for i in range(READER_TRAIN.num_docs)]), 1)\n",
    "            elif ARCH_TYPE == 1:    \n",
    "                out                 = torch.cat(tuple([net(None, torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][i]).to(DEVICE), torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][i]).to(DEVICE)) for i in range(READER_TRAIN.num_docs)]), 1)\n",
    "            else:\n",
    "                out                 = torch.cat(tuple([net(torch.from_numpy(features['local'][i]).to(DEVICE), torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][i]).to(DEVICE), torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][i]).to(DEVICE)) for i in range(READER_TRAIN.num_docs)]), 1)\n",
    "            loss                    = criterion(out, torch.from_numpy(features['labels']).to(DEVICE))\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            train_loss             += loss.item()\n",
    "        torch.save(net, MODEL_FILE.format(ens_idx + 1, ep_idx + 1))\n",
    "        print_message('model:{}, epoch:{}, loss:{}'.format(ens_idx + 1, ep_idx + 1, train_loss / EPOCH_SIZE))\n",
    "    is_complete                     = False\n",
    "    READER_DEV.reset()\n",
    "    net.eval()\n",
    "    while not is_complete:\n",
    "        features                    = READER_DEV.get_minibatch()\n",
    "        if ARCH_TYPE == 0:\n",
    "            out                     = net(torch.from_numpy(features['local'][0]).to(DEVICE), None, None)\n",
    "        elif ARCH_TYPE == 1:        \n",
    "            out                     = net(None, torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][0], torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][0]).to(DEVICE)).to(DEVICE))\n",
    "        else:\n",
    "            out                     = net(torch.from_numpy(features['local'][0]).to(DEVICE), torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][0]).to(DEVICE), torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][0]).to(DEVICE))\n",
    "        meta_cnt                    = len(features['meta'])\n",
    "        out                         = out.data.cpu()\n",
    "        for i in range(meta_cnt):\n",
    "            q                       = int(features['meta'][i][0])\n",
    "            d                       = int(features['meta'][i][1])\n",
    "            if q not in res_dev:    \n",
    "                res_dev[q]          = {}\n",
    "            if d not in res_dev[q]:\n",
    "                res_dev[q][d]       = 0\n",
    "            res_dev[q][d]          += out[i][0]\n",
    "        is_complete                 = (meta_cnt < MB_SIZE)\n",
    "    is_complete                     = False\n",
    "    READER_EVAL.reset()\n",
    "    net.eval()\n",
    "    while not is_complete:\n",
    "        features                    = READER_EVAL.get_minibatch()\n",
    "        if ARCH_TYPE == 0:\n",
    "            out                     = net(torch.from_numpy(features['local'][0]).to(DEVICE), None, None)\n",
    "        elif ARCH_TYPE == 1:\n",
    "            out                     = net(None, torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][0], torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][0]).to(DEVICE)).to(DEVICE))\n",
    "        else:\n",
    "            out                     = net(torch.from_numpy(features['local'][0]).to(DEVICE), torch.from_numpy(features['dist_q']).to(DEVICE), torch.from_numpy(features['dist_d'][0]).to(DEVICE), torch.from_numpy(features['mask_q']).to(DEVICE), torch.from_numpy(features['mask_d'][0]).to(DEVICE))\n",
    "        meta_cnt                    = len(features['meta'])\n",
    "        out                         = out.data.cpu()\n",
    "        for i in range(meta_cnt):\n",
    "            q                       = int(features['meta'][i][0])\n",
    "            d                       = int(features['meta'][i][1])\n",
    "            if q not in res_eval:\n",
    "                res_eval[q]         = {}\n",
    "            if d not in res_eval[q]:\n",
    "                res_eval[q][d]      = 0\n",
    "            res_eval[q][d]         += out[i][0]\n",
    "        is_complete                 = (meta_cnt < MB_SIZE)\n",
    "    mrr                                 = 0\n",
    "    for qid, docs in res_dev.items():\n",
    "        ranked                          = sorted(docs, key=docs.get, reverse=True)\n",
    "        for i in range(min(len(ranked), 10)):\n",
    "            if ranked[i] in qrels[qid]:\n",
    "                mrr                    += 1 / (i + 1)\n",
    "                break\n",
    "    mrr                                /= len(qrels)\n",
    "    print_message('model:{}, mrr:{}'.format(ens_idx + 1, mrr))\n",
    "with open(DATA_FILE_OUT_DEV, mode='w', encoding=\"utf-8\") as f:\n",
    "    for qid, docs in res_dev.items():\n",
    "        ranked                      = sorted(docs, key=docs.get, reverse=True)\n",
    "        for i in range(min(len(ranked), 10)):\n",
    "            f.write('{}\\t{}\\t{}\\n'.format(qid, ranked[i], i + 1))\n",
    "with open(DATA_FILE_OUT_EVAL, mode='w', encoding=\"utf-8\") as f:\n",
    "    for qid, docs in res_eval.items():\n",
    "        ranked                      = sorted(docs, key=docs.get, reverse=True)\n",
    "        for i in range(min(len(ranked), 10)):\n",
    "            f.write('{}\\t{}\\t{}\\n'.format(qid, ranked[i], i + 1))\n",
    "print_message('Finished')"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
